#ifndef TINYNDARRAY_INTERNAL_BINARY_H
#define TINYNDARRAY_INTERNAL_BINARY_H
#include <functional>
#include <sstream>
#include <type_traits>
#include <vector>

#include "broadcast.h"
#include "misc.h"

namespace tinyndarray {
// 包装二元操作用于指针接口
inline BinaryElementOp WrapBinaryOp(const BinaryElementOp& op) {
    return op;
}

// 通用的二元操作函数
template <typename F>
NdArray ApplyBinaryOp(const NdArray& lhs, const NdArray& rhs, F op) {
    if (lhs.shape() == rhs.shape()) {
        NdArray ret(lhs.dtype(), lhs.shape());
        ApplyOpSimple(ret, lhs, rhs, op);
        return ret;
    } else {
        Shape ret_shape = CheckBroadcastable(lhs.shape(), rhs.shape());
        NdArray ret(lhs.dtype(), ret_shape);
        ApplyOpBroadcast(ret, lhs, rhs, 0, 1, op);
        return ret;
    }
}

// 反转二元操作（用于标量在左侧的情况）
template <typename F>
struct ReverseAdapter {
    F func;
    ReverseAdapter(F f) : func(f) {}

    void operator()(void* out, const void* in1, const void* in2,
                    Dtype dtype) {
        func(out, in2, in1, dtype);  // 交换参数顺序
    }
};

// 数组与标量操作（泛型版本）
template <typename T>
auto ApplyBinaryOp(const NdArray& lhs, T scalar, ScalarElementOp op_func)
        -> std::enable_if_t<std::is_arithmetic_v<T>, NdArray> {
    NdArray scalar_arr(lhs.dtype(), {1});
    scalar_arr[{0}] = scalar;

    return ApplyBinaryOp(lhs, scalar_arr, op_func);
}

// 标量与数组操作（泛型版本）
template <typename T>
auto ApplyBinaryOp(T scalar, const NdArray& rhs, ScalarElementOp op_func)
        -> std::enable_if_t<std::is_arithmetic_v<T>, NdArray> {
    NdArray scalar_arr(rhs.dtype(), {1});
    scalar_arr[{0}] = scalar;
    // 应用反转操作
    ReverseAdapter<ScalarElementOp> reversed_op(op_func);

    return ApplyBinaryOp(scalar_arr, rhs, reversed_op);
}

}  // namespace tinyndarray
#endif